clc;
clear;
close all;

addpath('autogen\')

% 验证DH参数 passed
%{ 
a = [0, 0.062, 0.0];
alpha = [-pi/2, -pi/2, pi/2];

d = [0.0, 0.0, 0.0];
q = [-pi/2, -pi/2, 0];

for i=1:3
    T(:,:,i)=[
            cos(q(i)), -sin(q(i)), 0, a(i);
            sin(q(i)).*cos(alpha(i)), cos(q(i))*cos(alpha(i)), -sin(alpha(i)), -d(i)*sin(alpha(i));
            sin(q(i))*sin(alpha(i)), cos(q(i))*sin(alpha(i)), cos(alpha(i)), d(i)*cos(alpha(i));
            0,0,0,1];
end

P = [-0.11565; 0.0138; -0.92350; 1];
T01=T(:,:,1);
T02=T01*T(:,:,2);
T03=T02*T(:,:,3);
newP = inv(T03)*P;
%}

% 生成雅可比函数
%{
q=sym('q_%d',[3, 1],'real');
d = [0, 0, 0];
alpha = sym('alpha_%d',[3, 1],'real');
a = [0, 0.062, 0.0];

for i=1:3
    T(:,:,i)=[
            cos(q(i)), -sin(q(i)), 0, a(i);
            sin(q(i)).*cos(alpha(i)), cos(q(i))*cos(alpha(i)), -sin(alpha(i)), -d(i)*sin(alpha(i));
            sin(q(i))*sin(alpha(i)), cos(q(i))*sin(alpha(i)), cos(alpha(i)), d(i)*cos(alpha(i));
            0,0,0,1];
end

T01=T(:,:,1);
T02=T01*T(:,:,2);
T03=T02*T(:,:,3);

for (i=1: 4)
    for (j=1: 4)    
        % t1
        u11(i, j) = jacobian(T01(i, j), q(1));
        u21(i, j) = jacobian(T02(i, j), q(1));
        u31(i, j) = jacobian(T03(i, j), q(1));
        
        % t2
        u32(i, j) = jacobian(T03(i, j), q(2));
        u22(i, j) = jacobian(T02(i, j), q(2));
        
        % t3
        u33(i, j) = jacobian(T03(i, j), q(3));
    end
end

optimize = true;
matlabFunction(u11, 'File', 'autogen/J_h_T01_q1', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u21, 'File', 'autogen/J_h_T02_q1', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u31, 'File', 'autogen/J_h_T03_q1', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u32, 'File', 'autogen/J_h_T03_q2', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u22, 'File', 'autogen/J_h_T02_q2', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u33, 'File', 'autogen/J_h_T03_q3', 'Vars',{q, alpha}, 'Optimize', optimize);
%}

m1=0.13; m2=0.07; m3=0.5;
r1_x=0.05; r1_y=0.00; r1_z=0.0;
r2_x=0.0; r2_y=0.03; r2_z=0.0;
r3_x=0.0; r3_y=-0.02; r3_z=-0.08;

m1_r1 = m1*[r1_x;r1_y;r1_z;1];
m2_r2 = m2*[r2_x;r2_y;r2_z;1];
m3_r3 = m3*[r3_x;r3_y;r3_z;1];

lb = [0, 0, 0,...
    r1_x-0.0, r1_y-0.0, r1_z-0.0,...
    r2_x-0.0, r2_y-0.0, r2_z-0.0,...
    r3_x-0.1, r3_y-0.1, r3_z-0.5];

ub = [m1+0.1, m2+0.1, m3+3.0,...
    r1_x+0.0, r1_y+0.0, r1_z+0.0,...
    r2_x+0.0, r2_y+0.0, r2_z+0.0,...
    r3_x+0.1, r3_y+0.1, r3_z+0.1];

%{
ang_01 = rotm2eul(T01(1:3, 1:3), 'ZYX');
ang_02 = rotm2eul(T02(1:3, 1:3), 'ZYX');


q=sym('q_%d',[5, 1],'real');
d = [0, 0, 0.2765, 0, 0.1265];
alpha = sym('alpha_%d',[5, 1],'real');

%齐次变换矩阵
% SDH
% for i=1:5
% T(:,:,i)=[cos(q(i)),-sin(q(i))*cos(alpha(i)),sin(q(i))*sin(alpha(i)),a(i)*cos(q(i));
%         sin(q(i)),cos(q(i))*cos(alpha(i)),-cos(q(i))*sin(alpha(i)),a(i)*sin(q(i));
%         0,sin(alpha(i)),cos(alpha(i)),d(i);
%         0,0,0,1];
% end
% MDH
for i=1:5
    T(:,:,i)=[
            cos(q(i)), -sin(q(i)), 0, a(i);
            sin(q(i)).*cos(alpha(i)), cos(q(i))*cos(alpha(i)), -sin(alpha(i)), -d(i)*sin(alpha(i));
            sin(q(i))*sin(alpha(i)), cos(q(i))*sin(alpha(i)), cos(alpha(i)), d(i)*cos(alpha(i));
            0,0,0,1];
end

% 5连杆臂
T01=T(:,:,1);
T02=T01*T(:,:,2);
T03=T02*T(:,:,3);
T04=T03*T(:,:,4);
T05=T04*T(:,:,5);

for (i=1: 4)
    for (j=1: 4)
        % t1
        u51(i, j) = jacobian(T05(i, j), q(1));
        u41(i, j) = jacobian(T04(i, j), q(1));
        u31(i, j) = jacobian(T03(i, j), q(1));
        u21(i, j) = jacobian(T02(i, j), q(1));
        u11(i, j) = jacobian(T01(i, j), q(1));
        
        % t2
        u52(i, j) = jacobian(T05(i, j), q(2));
        u42(i, j) = jacobian(T04(i, j), q(2));
        u32(i, j) = jacobian(T03(i, j), q(2));
        u22(i, j) = jacobian(T02(i, j), q(2));
        
        % t3
        u53(i, j) = jacobian(T05(i, j), q(3));
        u43(i, j) = jacobian(T04(i, j), q(3));
        u33(i, j) = jacobian(T03(i, j), q(3));
        
        % t4
        u54(i, j) = jacobian(T05(i, j), q(4));
        u44(i, j) = jacobian(T04(i, j), q(4));
        
        % t5
        u55(i, j) = jacobian(T05(i, j), q(5));
    end
end

% 符号函数 -> matlab函数
optimize = true;
matlabFunction(u55, 'File', 'autogen/J_T05_q5', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u54, 'File', 'autogen/J_T05_q4', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u44, 'File', 'autogen/J_T04_q4', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u53, 'File', 'autogen/J_T05_q3', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u43, 'File', 'autogen/J_T04_q3', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u33, 'File', 'autogen/J_T03_q3', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u52, 'File', 'autogen/J_T05_q2', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u42, 'File', 'autogen/J_T04_q2', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u32, 'File', 'autogen/J_T03_q2', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u22, 'File', 'autogen/J_T02_q2', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u51, 'File', 'autogen/J_T05_q1', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u41, 'File', 'autogen/J_T04_q1', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u31, 'File', 'autogen/J_T03_q1', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u21, 'File', 'autogen/J_T02_q1', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u11, 'File', 'autogen/J_T01_q1', 'Vars',{q, alpha}, 'Optimize', optimize);
%}

%%%%%%%%%%%%%% 重力项标定 %%%%%%%%%%%%%%
% sota params
% m5_r5 = -[-0.0150033113572837;0.0155621404089697;0.00298782583024790;-0.0298183748107559];
% m4_r4 = -[0.0434142634867931;-0.00298782583024791;-0.00483493140326279;-0.0301963347782822];
% m3_r3 = -[0.00223801675774612;0.00483493140326279;-0.111130500367462;-0.0301963347782822];
% m2_r2 = -[0.0531248573321123;0.111130500367462;-0.00754865682560412;0];
% m1_r1 = -[0.0509152360511520;-0.00754865682560412;0;0];

% sota-2
% m5_r5 = -[0.00660547912526246;-0.00551136150340125;-0.0141660415106262;-0.0338427772057026];
% m4_r4 = -[-0.00770376908022287;0.0141660415106262;0.000484550892512696;-0.0320507729546084];
% m3_r3 = -[-0.0120626965934373;-0.000484550892512710;-0.116542153993433;-0.0320507729546084];
% m2_r2 = -[0.0173132624575914;0.116542153993433;0.00197770960705462;0];
% m1_r1 = -[-0.00604154500906833;0.00197770960705462;0;0];


% [-0.0238732860329777;0.00218499367896241;0;0;0.0181234442486960;0.113900565278430;0.00218499367896241;0;-0.0323468472598316;-0.00129301735289270;-0.113900565278430;-0.0313122718569988;0.00424734786405540;0.0250648116104718;0.00129301735289269;-0.0313122718569988;0.00223307762940797;-0.000816350489425599;-0.0250648116104718;-0.0344829705257235]
%std::vector<double> r_arm_pose = {2.6, -4, 2.8, -3.17, 1.57};
%std::cout << tau_est << "\n"; // -2.22153 -0.849529 0.253151 0.0297406 0
    
% 加载最新估计的参数
% load('est.mat');
% m5_r5 = -est(17:20);
% m4_r4 = -est(13:16);
% m3_r3 = -est(9:12);
% m2_r2 = -est(5:8);
% m1_r1 = -est(1:4);

% 读取数据
data = load('data_l3/head.txt');
% data = data(1:298, :);

alpha = [-pi/2, -pi/2, pi/2]';

qs = data(:, 1:3);
tor = data(:, 7:9);

g = [0 0 -9.8 0]';


K = [];
Tau = [];

tau_1s = [];tau_2s = [];tau_3s = [];
tau_1s_est = [];tau_2s_est = [];tau_3s_est = [];
for i=1:1:length(qs)
    q = qs(i, :)'; 
    
    u33 = J_h_T03_q3(q, alpha);
    u22 = J_h_T02_q2(q, alpha);u32 = J_h_T03_q2(q, alpha);
    u11 = J_h_T01_q1(q, alpha);u21 = J_h_T02_q1(q, alpha);u31 = J_h_T03_q1(q, alpha);
     
    tau_3s = [tau_3s; tor(i, 3)];
    tau_2s = [tau_2s; tor(i, 2)];
    tau_1s = [tau_1s; tor(i, 1)];
    
    Tau = [Tau; 
        tor(i, 1)'; 
        tor(i, 2)'; 
        tor(i, 3)'
        ];
    
    tau_3 = -g'*u33*m3_r3;
    tau_2 = -g'*u22*m2_r2-g'*u32*m3_r3;
    tau_1 = -g'*u11*m1_r1-g'*u21*m2_r2-g'*u31*m3_r3;
    
    tau_3s_est = [tau_3s_est; tau_3];
    tau_2s_est = [tau_2s_est; tau_2];
    tau_1s_est = [tau_1s_est; tau_1];
end

% figure(1);
% plot(tau_1s, "r.-");
% hold on;
% plot(tau_1s_est, "g.-");
% legend('real', 'est');   
% figure(2);
% plot(tau_2s, "r.-");
% hold on;
% plot(tau_2s_est, "g.-");
% legend('real', 'est');   
% figure(3);
% plot(tau_3s, "r.-");
% hold on;
% plot(tau_3s_est, "g.-");
% legend('real', 'est');   

% 非线性最小二乘估计
x_init = [
    m1, m2, m3, ...
    r1_x, r1_y, r1_z, ...
    r2_x, r2_y, r2_z, ...
    r3_x, r3_y, r3_z];

qs_input = qs;
Taus = tau_estimator(x_init, qs);

x = lsqcurvefit(@tau_estimator, x_init, qs_input, Tau);
% [0.0355400542392994;0.0181187668264032;0;0.207278151718490]
% [-0.0888652885381021;0.0164045541206430;0.0117166062203967;0.281071516982919]
% [0.0136544164046772;-0.00728086706865319;-0.157040084611476;0.578161012781261]
% 
% [0.0427912005847312;0.0222120830358785;0;0.224000372144442]
% [-0.0774110496186375;0.0149538697111237;0.0133614350738391;0.265692916282782]
% [0.0210219346827239;-0.0124369973393612;-0.148797567090943;0.588365793101965]

% x = lsqcurvefit(@tau_estimator, x_init, qs_input, Tau, lb, ub);

m1_r1 = [x(1)*x(4); x(1)*x(5); x(1)*x(6); x(1)]; % [0.0101411505636156;0;0;0.202823011272312]
m2_r2 = [x(2)*x(7); x(2)*x(8); x(2)*x(9); x(2)]; % [0;0.00410734266715084;0;0.136911422238361]
m3_r3 = [x(3)*x(10); x(3)*x(11); x(3)*x(12); x(3)];% [-0.0643570858037036;0.0132101808416077;-0.150433751369505;1.38273061049691]
    
tau_1s_real = [];tau_2s_real = [];tau_3s_real = [];
tau_1s_est = [];tau_2s_est = [];tau_3s_est = [];
for i=1:1:length(qs_input)
    q = qs_input(i, :)'; 

    u33 = J_h_T03_q3(q, alpha);
    u22 = J_h_T02_q2(q, alpha);u32 = J_h_T03_q2(q, alpha);
    u11 = J_h_T01_q1(q, alpha);u21 = J_h_T02_q1(q, alpha);u31 = J_h_T03_q1(q, alpha);

    tau_3 = -g'*u33*m3_r3;
    tau_2 = -g'*u22*m2_r2-g'*u32*m3_r3;
    tau_1 = -g'*u11*m1_r1-g'*u21*m2_r2-g'*u31*m3_r3;

    tau_3s_real = [tau_3s_real; tor(i, 3)];
    tau_2s_real = [tau_2s_real; tor(i, 2)];
    tau_1s_real = [tau_1s_real; tor(i, 1)];
    
    tau_3s_est = [tau_3s_est; tau_3];
    tau_2s_est = [tau_2s_est; tau_2];
    tau_1s_est = [tau_1s_est; tau_1];
end


% plot
figure(1);
plot(tau_1s_real, "r.-");
hold on;
plot(tau_1s_est, "g.-");
figure(2);
plot(tau_2s_real, "r.-");
hold on;
plot(tau_2s_est, "g.-");
figure(3);
plot(tau_3s_real, "r.-");
hold on;
plot(tau_3s_est, "g.-");
legend('real', 'est');   

function taus = tau_estimator(x, qs_input)
    alpha = [-pi/2, -pi/2, pi/2]';
    g = [0 0 -9.8 0]';
    
    m1_r1 = [x(1)*x(4); x(1)*x(5); x(1)*x(6); x(1)];
    m2_r2 = [x(2)*x(7); x(2)*x(8); x(2)*x(9); x(2)];
    m3_r3 = [x(3)*x(10); x(3)*x(11); x(3)*x(12); x(3)];
    
    taus = [];
    for i=1:1:length(qs_input)
        q = qs_input(i, :)'; 

        u33 = J_h_T03_q3(q, alpha);
        u22 = J_h_T02_q2(q, alpha);u32 = J_h_T03_q2(q, alpha);
        u11 = J_h_T01_q1(q, alpha);u21 = J_h_T02_q1(q, alpha);u31 = J_h_T03_q1(q, alpha);

        tau_3 = -g'*u33*m3_r3;
        tau_2 = -g'*u22*m2_r2-g'*u32*m3_r3;
        tau_1 = -g'*u11*m1_r1-g'*u21*m2_r2-g'*u31*m3_r3;
        
        taus = [
            taus; 
            tau_1; 
            tau_2; 
            tau_3; 
            ];
    end
end
